package outputparserx

import (
	"encoding/json"
	"fmt"
	"github.com/tmc/langchaingo/outputparser"
	"strings"
)

// ParseError is the error type returned by output parsers.
type ParseError struct {
	Text   string
	Reason string
}

func (e ParseError) Error() string {
	return fmt.Sprintf("parse text %s. %s", e.Text, e.Reason)
}

const (
	// _structuredFormatInstructionTemplate is a template for the format
	// instructions of the structured output parser.
	_structuredFormatInstructionTemplate = "The output should be a markdown code snippet formatted in the following schema: \n```json\n%s\n```" // nolint

	// _structuredLineTemplate is a single line of the json schema in the
	// format instruction of the structured output parser. The fist verb is
	// the name, the second verb is the type and the third is a description of
	// what the field should contain.
	_structuredLineTemplate = "\"%s\": %s // %s\n"
)

type ResponseSchema struct {
	Name        string
	Description string
	Types       string
	Require     bool // 是否必选
	Schemas     []ResponseSchema
}

type Structured struct {
	outputparser.Structured
	schema []ResponseSchema
}

func NewStructured(schema []ResponseSchema) Structured {
	oResponseSchema := make([]outputparser.ResponseSchema, 0, len(schema))
	for i, _ := range schema {
		oResponseSchema = append(oResponseSchema, outputparser.ResponseSchema{
			Name:        schema[i].Name,
			Description: schema[i].Description,
		})
	}
	return Structured{
		schema:     schema,
		Structured: outputparser.NewStructured(oResponseSchema),
	}
}

func (p Structured) Parse(text string) (any, error) {
	return p.parse(text)
}

func (p Structured) parse(text string) (map[string]any, error) {

	var parsed map[string]any

	if err := json.Unmarshal([]byte(text), &parsed); err != nil {
		// Remove the ```json that should be at the start of the text, and the ```
		// that should be at the end of the text.
		withoutJSONStart := strings.Split(text, "```json")
		if !(len(withoutJSONStart) > 1) {
			return nil, ParseError{Text: text, Reason: "no ```json at start of output"}
		}

		withoutJSONEnd := strings.Split(withoutJSONStart[1], "```")
		if len(withoutJSONEnd) < 1 {
			return nil, ParseError{Text: text, Reason: "no ``` at end of output"}
		}

		jsonString := withoutJSONEnd[0]
		err := json.Unmarshal([]byte(jsonString), &parsed)
		if err != nil {
			return nil, err
		}
	}

	// Validate that the parsed map contains all fields specified in the response
	// schemas.
	missingKeys := make([]string, 0)
	for _, rs := range p.schema {
		if _, ok := parsed[rs.Name]; !ok && rs.Require {
			missingKeys = append(missingKeys, rs.Name)
		}
	}

	if len(missingKeys) > 0 {
		return nil, ParseError{
			Text:   text,
			Reason: fmt.Sprintf("output is missing the following fields %v", missingKeys),
		}
	}

	return parsed, nil
}

func (p Structured) GetFormatInstructions() string {
	//jsonLines := ""
	//for _, rs := range p.schema {
	//	if len(rs.Types) == 0 {
	//		rs.Types = "string"
	//	}
	//	jsonLines += "\t" + fmt.Sprintf(
	//		_structuredLineTemplate,
	//		rs.Name,
	//		rs.Types,
	//		rs.Description,
	//	)
	//}

	return fmt.Sprintf(_structuredFormatInstructionTemplate, p.jsonMarshal(p.schema, 0))
}

func (p Structured) jsonMarshal(schemas []ResponseSchema, level int) string {

	level++

	endBlank := ""
	fieldBlank := "\t"
	for i := 1; i < level; i++ {
		fieldBlank += "\t"
		endBlank += "\t"
	}

	jsonLines := "{"
	for _, rs := range schemas {
		if len(rs.Schemas) > 0 {
			rs.Types = p.jsonMarshal(rs.Schemas, level)
		}

		blank := fieldBlank
		if len(jsonLines) == 1 {
			blank = "\t"
		}

		if len(rs.Types) == 0 {
			rs.Types = "string"
		}
		jsonLines += blank + fmt.Sprintf(
			_structuredLineTemplate,
			rs.Name,
			rs.Types,
			rs.Description,
		)
	}

	return jsonLines + endBlank + "}"
}
